# -*- coding: utf-8 -*-
"""
Created on Sat Jun 27 12:28:46 2020

@author: aaa
"""


import matplotlib.pyplot as plt
import numpy as np
import time
import os
import tensorflow as tf

def unpickle(file_path):
    import pickle
    with open(file_path,'rb') as fo:
        dict = pickle.load(fo,encoding = 'bytes')
    return dict
data_path = []

def data_load():
    data_totle = []
    labels = []
    labels_name = []
    for i in range(5):
        k = i + 1
        data_path.append('E:\\文件\\vision\\cifar-10-batches-py\\data_batch_' + str(k))
        data_batch = unpickle(data_path[-1])
        key = data_batch.keys()
        value = list(data_batch.values())
        
        images = value[2]
        data_totle += list(images)
        labels += value[1]
        labels_name += value[3]
    return np.array(data_totle),np.array(labels),np.array(labels_name)

images,labels,labels_name = data_load()
print('图片尺寸：',images.shape)  #50000,3072
print(labels.shape,  type(images))

def l01(x):
    a = (x - np.min(x,axis = 1).reshape([-1,1]))
    b = (np.max(x,axis = 1) - np.min(x,axis = 1)).reshape([-1,1])
    return (a / b)
images = l01(images) 

def one_hot(x):
    x_hot = np.zeros((len(x),10))
    x_hot[np.arange(len(x)),x] = 1
    return x_hot
labels = one_hot(labels)
print(labels[1])
print(labels_name[1])

#模型参数
lr = 1e-5
n_features = 3072
n_w1 = 32
n_w2 = 64
n_fc1 = 384
n_fc2 = 192
n_class = 10


#构建模型，初始化
x = tf.placeholder(dtype = tf.float32,shape = [None,n_features])
y = tf.placeholder(dtype = tf.float32,shape = [None,n_class])

def w_init(shape,stddev = 0.1):
    return tf.Variable(tf.truncated_normal(shape = shape,stddev = stddev))
def b_init(shape,value = 0.1):
    return tf.Variable(tf.constant(value,dtype = tf.float32,shape = shape ))

w1 = w_init([5,5,3,32],stddev = 0.0001)
b1 = b_init([32],value = 0.0)
w2 = w_init([5,5,32,64],stddev = 0.01)
b2 = b_init([64])

w_fc1 = w_init([8 * 8 * 64,n_fc1])
b_fc1 = b_init([n_fc1])
w_fc2 = w_init([n_fc1,n_fc2])
b_fc2 = b_init([n_fc2])
w_fc3 = w_init([n_fc2,n_class])
b_fc3 = b_init([n_class],value = 0.0)


x_image = tf.reshape(x,[-1,32,32,3])
#定义卷积，池化，relu，
def conv(x,w):
    return tf.nn.conv2d(x,w,strides = [1,1,1,1],padding = 'SAME')
def relu(x):
    return tf.nn.relu(x)
def maxpool(x):
    return tf.nn.max_pool(x,ksize = [1,3,3,1],strides = [1,2,2,1],padding = 'SAME')
def softmax(x):
    return tf.nn.softmax(x)


#conv1
c1 = conv(x_image,w1)
h_c1 = relu(c1 + b1)
pool1 = maxpool(h_c1)

#conv2
c2 = conv(pool1,w2)
h_c2 = relu(c2 + b2)
pool2 = maxpool(h_c2)

#fc1 
flat = tf.reshape(pool2,[-1,8 * 8 * 64])
fc1 = tf.matmul(flat,w_fc1) + b_fc1
h_fc1 = relu(fc1)

#fc2 
fc2 = tf.matmul(h_fc1,w_fc2) + b_fc2
h_fc2 = relu(fc2)

#fc2-softmax
fc3 = tf.matmul(h_fc2,w_fc3) + b_fc3
output = softmax(fc3)

#损失函数loss，梯度更新
loss = tf.reduce_mean(-tf.reduce_sum(y * tf.log(output)))
#train = tf.train.GradientDescentOptimizer(lr).minimize(loss)
train = tf.train.AdamOptimizer(learning_rate= lr).minimize(loss)

saver = tf.train.Saver()
if not os.path.exists('tmp/'):
    os.mkdir('tmp/')

#创建会话
with tf.Session() as sess:
    if os.path.exists('tmp/checkpoint'):          #判断模型是否存在
        saver.restore(sess,'tmp/model.ckpt')      #从模型中恢复变量
    else:
        sess.run(tf.global_variables_initializer())
    num = len(labels)
    #num = 30000
    training_iters = 100
    batch_size = 200
    start_time = time.time()
    cost =[]
    t = 0
    m = 2000                                        #进度间隔
    for epoch in range(training_iters):
        print('第 %d 次迭代' % (epoch + 1))
        for batch in range(0,num,batch_size):
            
            x_batch = images[batch:batch + batch_size,:]
            y_batch = labels[batch:batch + batch_size,:]
            
            probility = sess.run(output,feed_dict = {x:x_batch,y:y_batch})
            accuracy = np.sum(np.argmax(probility,axis = 1) == np.argmax(y_batch,axis = 1)) / batch_size
            
            t += 1
            sess.run(train,feed_dict = {x:x_batch,y:y_batch})
            if lr > 1e-5:
                if batch % 5000 == 0:
                    lr = lr / 5
            if batch % m == 0:
                j = '进度:' + str(format(  ((batch + m) / num) * 100,'.2f'  ) ) + '%'
                print('\r' + j,end = '')
                time.sleep(0.3)
                
        
        cost.append(sess.run(loss,feed_dict = {x:x_batch ,y:y_batch}))
        print()
        print('损失值:',cost[-1])
        print('精度为:%.3f,  学习率为：%.8f' % (accuracy,lr))
        print()
        
    save_paramters = saver.save(sess,'tmp/model.ckpt')    #写入当前模型中所有可训练变量 
    end_time = time.time()
    
    plt.plot(cost)
    title = 'learning_rate %f, cost %f' % (lr,cost[-1])
    plt.title(title)
    plt.ylabel('cost')
    plt.show()
    print(cost[0],   cost[-1])
    print('运行时间：%.3f 分钟' % ((end_time - start_time) / 60))
    
